// Microsoft Public License (Ms-PL)
// Copyright (c) 2010 Nick Darnell
//
// http://www.nickdarnell.com
// http://wpfshadergenerator.codeplex.com

using System;
using System.CodeDom;
using System.CodeDom.Compiler;
using System.Diagnostics;
using System.IO;
using System.Runtime.InteropServices;
using System.Text;
using System.Text.RegularExpressions;
using System.Windows;
using System.Windows.Media;
using System.Windows.Media.Effects;
using System.Windows.Media.Media3D;
using System.Xml;
using System.Xml.Schema;
using Microsoft.VisualStudio.Shell;
using Microsoft.Win32;
using Microsoft.VisualStudio.Designer.Interfaces;
using Microsoft.VisualStudio.Shell.Interop;

namespace WPFShaderEffectGenerator
{
    [ComVisible(true)]
    [Guid("1F271326-3217-40f7-B1B9-6D88891532C2")]
    [ProvideObject(typeof(ShaderEffectGenerator), RegisterUsing = RegistrationMethod.CodeBase)]
    [CodeGeneratorRegistration(typeof(ShaderEffectGenerator), "ShaderEffectGenerator", EnvDTEConstants.vsContextGuidVBProject, GeneratesDesignTimeSource = true)]
    [CodeGeneratorRegistration(typeof(ShaderEffectGenerator), "ShaderEffectGenerator", EnvDTEConstants.vsContextGuidVCSProject, GeneratesDesignTimeSource = true)]
    public class ShaderEffectGenerator : BaseCodeGeneratorWithSite
    {
        protected override byte[] GenerateCode(string inputFileContent)
        {
            ShaderFragmentReader.LogError = GeneratorError;
            ShaderFragmentReader.LogWarning = GeneratorWarning;

            try
            {
                // Try to generate the wrapper file.
                return GenerateCodeSafe(inputFileContent);
            }
            catch (Exception ex)
            {
                // In case of a faliure - print the exception 
                // as a comment on the .cs file.
                return GenerateExceptionMessage(ex);
            }
        }

        private byte[] GenerateCodeSafe(string inputFileContent)
        {
            CodeCompileUnit code = new CodeCompileUnit();

            // Parse Fragement
            ShaderFragment fragement = ShaderFragmentReader.Read(base.InputFilePath, inputFileContent);

            // Create Namespce
            CodeNamespace codeNamespace = CreateNamespace(code);

            // Compile pixel shader.
            CompilePixelShader(inputFileContent, fragement);

            // Create the ShaderEffect Class
            CreateEffectClass(codeNamespace, base.InputFilePath, fragement);

            string stringcode = WriteCode(GetCodeProvider(), code);

            return Encoding.ASCII.GetBytes(stringcode);
        }

        private CodeNamespace CreateNamespace(CodeCompileUnit code)
        {
            CodeNamespace codeNamespace = new CodeNamespace(FileNameSpace);
            code.Namespaces.Add(codeNamespace);

            codeNamespace.Comments.Add(new CodeCommentStatement("-----------------------------------------------------------------------------"));
            codeNamespace.Comments.Add(new CodeCommentStatement(" <auto-generated>"));
            codeNamespace.Comments.Add(new CodeCommentStatement("     This code was generated by:"));
            codeNamespace.Comments.Add(new CodeCommentStatement("       The WPF ShaderEffect Generator"));
            codeNamespace.Comments.Add(new CodeCommentStatement("       http://wpfshadergenerator.codeplex.com"));
            codeNamespace.Comments.Add(new CodeCommentStatement(""));
            codeNamespace.Comments.Add(new CodeCommentStatement("     Changes to this file may cause incorrect behavior and will be lost if"));
            codeNamespace.Comments.Add(new CodeCommentStatement("     the code is regenerated."));
            codeNamespace.Comments.Add(new CodeCommentStatement(" </auto-generated>"));
            codeNamespace.Comments.Add(new CodeCommentStatement("-----------------------------------------------------------------------------"));


            codeNamespace.Imports.Add(new CodeNamespaceImport("System"));
            codeNamespace.Imports.Add(new CodeNamespaceImport("System.Windows"));
            codeNamespace.Imports.Add(new CodeNamespaceImport("System.Windows.Media"));
            codeNamespace.Imports.Add(new CodeNamespaceImport("System.Windows.Media.Effects"));
            codeNamespace.Imports.Add(new CodeNamespaceImport("System.Windows.Media.Media3D"));

            return codeNamespace;
        }

        String CSharp_RegisterPixelShaderSampler =
            "        public static readonly DependencyProperty {0}Property = ShaderEffect.RegisterPixelShaderSamplerProperty(\"{0}\", typeof({1}), {2}, {3});";

        String CSharp_RegisterDependencyProperty =
            "        public static readonly DependencyProperty {0}Property = DependencyProperty.Register(\"{0}\", typeof({3}), typeof({1}), new PropertyMetadata({4}, PixelShaderConstantCallback({2})));";

        private string DefaultConverter(object value)
        {
            Type valueType = value.GetType();

            if (valueType == typeof(double))
            {
                return "(double)" + value.ToString();
            }
            else if (valueType == typeof(float))
            {
                return "(float)" + value.ToString();
            }
            else if (valueType == typeof(Point))
            {
                Point p = (Point)value;
                return "new Point(" + p.X + "," + p.Y + ")";
            }
            else if (valueType == typeof(Size))
            {
                Size s = (Size)value;
                return "new Size(" + s.Width + "," + s.Height + ")";
            }
            else if (valueType == typeof(Vector))
            {
                Vector v = (Vector)value;
                return "new Vector(" + v.X + "," + v.Y + ")";
            }
            else if (valueType == typeof(Point3D))
            {
                Point3D p = (Point3D)value;
                return "new Point3D(" + p.X + "," + p.Y + "," + p.Z + ")";
            }
            else if (valueType == typeof(Vector3D))
            {
                Vector3D v = (Vector3D)value;
                return "new Vector3D(" + v.X + "," + v.Y + "," + v.Z + ")";
            }
            else if (valueType == typeof(Point4D))
            {
                Point4D p = (Point4D)value;
                return "new Point4D(" + p.X + "," + p.Y + "," + p.Z + "," + p.W + ")";
            }
            else if (valueType == typeof(Color))
            {
                Color c = (Color)value;
                return "Color.FromArgb(" + c.A + "," + c.R + "," + c.G + "," + c.B + ")";
            }
            else if (valueType == typeof(bool))
            {
                return value.ToString();
            }
            else if (valueType == typeof(int))
            {
                return "(int)" + value.ToString();
            }
            else if (valueType == typeof(uint))
            {
                return "(uint)" + value.ToString();
            }
            else if (valueType == typeof(byte))
            {
                return "(byte)" + value.ToString();
            }
            else if (valueType == typeof(sbyte))
            {
                return "(sbyte)" + value.ToString();
            }
            else if (valueType == typeof(short))
            {
                return "(short)" + value.ToString();
            }
            else if (valueType == typeof(ushort))
            {
                return "(ushort)" + value.ToString();
            }
            else if (valueType == typeof(char))
            {
                return "(char)" + ((int)(char)value).ToString();
            }
            else if (valueType == typeof(long))
            {
                return "(long)" + value.ToString();
            }
            else if (valueType == typeof(ulong))
            {
                return "(ulong)" + value.ToString();
            }
            else
            {
                return "null";
            }
        }

        private CodeTypeDeclaration CreateEffectClass(CodeNamespace codeNamespace, string inputFileName, ShaderFragment fragement)
        {
            // Create the header class
            string pixelShaderFileName = Path.GetFileNameWithoutExtension(inputFileName) + ".ps";
            string ownerClass = Path.GetFileNameWithoutExtension(inputFileName);
            CodeTypeDeclaration wrapperClass = new CodeTypeDeclaration(ownerClass);
            wrapperClass.BaseTypes.Add(new CodeTypeReference(typeof(ShaderEffect)));
            wrapperClass.Comments.Add(new CodeCommentStatement(new CodeComment("<summary>" + fragement.Description + "</summary>", true)));

            foreach (ShaderFragmentRegister register in fragement.Registers)
            {
                CodeSnippetTypeMember memberField = null;
                if (register.Type == typeof(Brush))
                {
                    memberField = new CodeSnippetTypeMember(String.Format(CSharp_RegisterPixelShaderSampler, register.Name, ownerClass, register.Index, "SamplingMode." + register.SamplingMode));
                }
                else
                {
                    memberField = new CodeSnippetTypeMember(String.Format(CSharp_RegisterDependencyProperty, register.Name, ownerClass, register.Index, register.Type, DefaultConverter(register.DefaultValue)));
                }

                if (memberField != null)
                {
                    memberField.Comments.Add(new CodeCommentStatement(new CodeComment("<summary>" + register.Description + "</summary>", true)));
                    wrapperClass.Members.Add(memberField);
                }
            }

            foreach (ShaderFragmentRegister register in fragement.Registers)
            {
                CodeMemberProperty property = new CodeMemberProperty();
                property.Name = register.Name;
                property.Type = new CodeTypeReference(register.Type);
                property.Attributes = MemberAttributes.Public;
                property.Comments.Add(new CodeCommentStatement(new CodeComment("<summary>" + register.Description + "</summary>", true)));
                property.HasGet = true;
                property.HasSet = true;

                property.GetStatements.Add(
                    new CodeMethodReturnStatement(
                        new CodeCastExpression(register.Type,
                            new CodeMethodInvokeExpression(
                                new CodeMethodReferenceExpression(
                                    new CodeThisReferenceExpression(),
                                    "GetValue"),
                                new CodeFieldReferenceExpression(null, register.Name + "Property")))));

                property.SetStatements.Add(
                    new CodeMethodInvokeExpression(
                        new CodeMethodReferenceExpression(
                            new CodeThisReferenceExpression(),
                            "SetValue"),
                        new CodeFieldReferenceExpression(null, register.Name + "Property"),
                        new CodeSnippetExpression("value")));

                wrapperClass.Members.Add(property);
            }

            EnvDTE.ProjectItem item = base.GetService(typeof(EnvDTE.ProjectItem)) as EnvDTE.ProjectItem;
            EnvDTE.Project project = item.ContainingProject;

            string AssemblyName = project.Properties.Item("AssemblyName").Value.ToString();

            CodeConstructor constructor = new CodeConstructor();
            constructor.Name = ownerClass;
            constructor.Attributes = MemberAttributes.Public;

            string finalShaderPath = Path.GetDirectoryName(inputFileName) + "\\" + pixelShaderFileName;
            Uri projUri = new Uri(project.FileName);
            Uri psUri = projUri.MakeRelativeUri(new Uri(finalShaderPath));

            constructor.Statements.Add(new CodeSnippetExpression("PixelShader pixelShader = new PixelShader()"));
            constructor.Statements.Add(new CodeSnippetExpression("pixelShader.UriSource = new Uri(\"" +
                "/" + AssemblyName + ";component/" + psUri + "\", UriKind.Relative)"));
            constructor.Statements.Add(new CodeSnippetExpression("this.PixelShader = pixelShader"));

            foreach (ShaderFragmentRegister register in fragement.Registers)
            {
                constructor.Statements.Add(new CodeSnippetExpression("this.UpdateShaderValue(" + register.Name + "Property)"));
            }

            constructor.Statements.Add(new CodeSnippetExpression("this.DdxUvDdyUvRegisterIndex = " + fragement.DdxUvDdyUvRegisterIndex));

            wrapperClass.Members.Add(constructor);

            codeNamespace.Types.Add(wrapperClass);

            return wrapperClass;
        }

        private EnvDTE.ProjectItem CompilePixelShader(string inputFileContent, ShaderFragment fragement)
        {
            string finalShaderPath = Path.GetDirectoryName(base.InputFilePath) + "\\" + Path.GetFileNameWithoutExtension(base.InputFilePath) + ".ps";

            dx9fxutil.PixelShaderCompiler shaderCompiler = new dx9fxutil.PixelShaderCompiler();
            shaderCompiler.Compile(
                Path.GetFileName(base.InputFilePath),
                inputFileContent,
                String.IsNullOrEmpty(fragement.Profile) ? "ps_2_0" : fragement.Profile,
                finalShaderPath,
                GeneratorWarning,
                GeneratorError);

            if (File.Exists(finalShaderPath))
            {
                EnvDTE.ProjectItem item = base.GetService(typeof(EnvDTE.ProjectItem)) as EnvDTE.ProjectItem;
                EnvDTE.Project project = item.ContainingProject;

                EnvDTE.ProjectItem finalShaderItem =
                    item.ProjectItems.AddFromFile(finalShaderPath);
                finalShaderItem.Properties.Item("ItemType").Value = "Resource";

                return finalShaderItem;
            }

            return null;
        }

        private byte[] GenerateExceptionMessage(Exception ex)
        {
            CodeCompileUnit code = new CodeCompileUnit();
            CodeNamespace codeNamespace = new CodeNamespace(FileNameSpace);
            code.Namespaces.Add(codeNamespace);

            codeNamespace.Comments.Add(new CodeCommentStatement(ex.ToString()));

            string stringcode = WriteCode(base.GetCodeProvider(), code);

            return Encoding.ASCII.GetBytes(stringcode);
        }

        private static string WriteCode(CodeDomProvider p_provider, CodeCompileUnit code)
        {
            CodeGeneratorOptions options = new CodeGeneratorOptions();
            options.BlankLinesBetweenMembers = true;
            options.BracingStyle = "C";
            options.IndentString = "    ";

            StringWriter stringwriter = new StringWriter();
            p_provider.GenerateCodeFromCompileUnit(code, stringwriter, options);
            string stringcode = stringwriter.ToString();
            int index = stringcode.IndexOf("</auto-generated>");
            stringcode = stringcode.Remove(0, index + 103);
            stringwriter.Close();

            return stringcode;
        }

        private TargetFramework TargetFramework
        {
            get
            {
                EnvDTE.ProjectItem item = base.GetService(typeof(EnvDTE.ProjectItem)) as EnvDTE.ProjectItem;
                EnvDTE.Project project = item.ContainingProject;

                return (TargetFramework)(uint)project.Properties.Item("TargetFramework").Value;
            }
        }
    }
}